#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""修复 vector_store_id，确保其等于 document_id"""

import os
import sys

# Windows 控制台编码设置
if sys.platform == 'win32':
    import codecs
    sys.stdout = codecs.getwriter('utf-8')(sys.stdout.buffer, 'strict')
    sys.stderr = codecs.getwriter('utf-8')(sys.stderr.buffer, 'strict')

# 添加项目根目录到路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from app.models.database import engine
from sqlalchemy.orm import sessionmaker
from app.models.document import Document

def main():
    Session = sessionmaker(bind=engine)
    db = Session()
    
    try:
        docs = db.query(Document).filter(Document.is_vectorized == True).all()
        
        print(f'\n=== 修复 vector_store_id ===')
        fixed_count = 0
        
        for d in docs:
            if d.vector_store_id != d.id:
                print(f'修复: {d.title[:40]}')
                print(f'  旧 vector_store_id: {d.vector_store_id}')
                print(f'  新 vector_store_id: {d.id}')
                d.vector_store_id = d.id
                fixed_count += 1
        
        if fixed_count > 0:
            db.commit()
            print(f'\n✅ 成功修复 {fixed_count} 个文档的 vector_store_id')
        else:
            print(f'\n✅ 所有文档的 vector_store_id 都正确，无需修复')
        
    except Exception as e:
        db.rollback()
        print(f'\n❌ 修复失败: {str(e)}')
    finally:
        db.close()

if __name__ == "__main__":
    main()

